#!/usr/bin/env python3
import os
import signal
import logging
import json
from collections import defaultdict
"""
支持多线程和网络结构化的MCP服务器
实现絮状网络结构和环境隔离层
"""
import json
import sys
import os
import signal
import threading
import time
import queue
import asyncio
import concurrent.futures
import importlib
import inspect
import logging
import psutil
import socket
from typing import Dict, List, Any, Optional, Callable
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] %(levelname)s: %(message)s',
    stream=sys.stderr
)
logger = logging.getLogger('NetworkMCP')

@dataclass
class Environment:
    """环境配置"""
    name: str
    path: str
    modules: List[str]
    variables: Dict[str, str]

class EnvironmentManager:
    """环境管理器 - 处理不同执行环境"""
    
    def __init__(self):
        self.environments: Dict[str, Environment] = {}
        self.env_locks: Dict[str, threading.Lock] = {}
        self._discover_environments()
    
    def _discover_environments(self):
        """发现并注册可用的执行环境"""
        # 检测Python环境
        self.register_environment(
            "python",
            sys.prefix,
            ["json", "asyncio", "concurrent.futures"],
            {"PYTHONPATH": sys.prefix}
        )
        
        # 检测Node.js环境（如果可用）
        try:
            import subprocess
            result = subprocess.run(["node", "--version"], 
                                 capture_output=True, text=True)
            if result.returncode == 0:
                node_path = subprocess.run(["which", "node"], 
                                        capture_output=True, 
                                        text=True).stdout.strip()
                self.register_environment(
                    "node",
                    os.path.dirname(node_path),
                    ["npm", "node"],
                    {"NODE_PATH": os.path.dirname(node_path)}
                )
        except Exception as e:
            logger.debug(f"Node.js环境不可用: {e}")
    
    def register_environment(self, name: str, path: str, 
                           modules: List[str], variables: Dict[str, str]):
        """注册新的执行环境"""
        self.environments[name] = Environment(name, path, modules, variables)
        self.env_locks[name] = threading.Lock()
        logger.info(f"注册环境: {name} @ {path}")
    
    @contextmanager
    def use_environment(self, name: str):
        """使用特定环境的上下文管理器"""
        if name not in self.environments:
            raise ValueError(f"未知环境: {name}")
            
        env = self.environments[name]
        lock = self.env_locks[name]
        
        with lock:
            # 保存当前环境变量
            old_env = {}
            for key, value in env.variables.items():
                old_env[key] = os.environ.get(key)
                os.environ[key] = value
            
            try:
                yield env
            finally:
                # 恢复环境变量
                for key, value in old_env.items():
                    if value is None:
                        del os.environ[key]
                    else:
                        os.environ[key] = value

class MessageBroker:
    """消息代理 - 处理组件间通信"""
    
    def __init__(self):
        self.topics: Dict[str, queue.Queue] = {}
        self.subscribers: Dict[str, List[Callable]] = {}
        self._lock = threading.Lock()
    
    def create_topic(self, topic: str) -> None:
        """创建新的主题"""
        with self._lock:
            if topic not in self.topics:
                self.topics[topic] = queue.Queue()
                self.subscribers[topic] = []
                logger.info(f"创建主题: {topic}")
    
    def subscribe(self, topic: str, callback: Callable) -> None:
        """订阅主题"""
        with self._lock:
            if topic not in self.topics:
                self.create_topic(topic)
            self.subscribers[topic].append(callback)
            logger.debug(f"订阅主题 {topic}: {callback.__name__}")
    
    def publish(self, topic: str, message: Any) -> None:
        """发布消息到主题"""
        with self._lock:
            if topic not in self.topics:
                self.create_topic(topic)
            
            # 将消息放入队列
            self.topics[topic].put(message)
            
            # 通知所有订阅者
            for callback in self.subscribers[topic]:
                try:
                    callback(message)
                except Exception as e:
                    logger.error(f"处理消息时出错: {e}")

class RequestRouter:
    """请求路由器 - 处理请求分发"""
    
    def __init__(self, thread_pool: ThreadPoolExecutor):
        self.routes: Dict[str, Callable] = {}
        self.thread_pool = thread_pool
        self._lock = threading.Lock()
    
    def register_route(self, method: str, handler: Callable) -> None:
        """注册新路由"""
        with self._lock:
            self.routes[method] = handler
            logger.info(f"注册路由: {method}")
    
    async def route_request(self, method: str, params: Dict) -> Any:
        """路由并处理请求"""
        if method not in self.routes:
            raise ValueError(f"未知方法: {method}")
            
        handler = self.routes[method]
        
        # 在线程池中执行处理程序
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            self.thread_pool,
            handler,
            params
        )

class NetworkServer:
    """网络服务器实现"""
    
    def __init__(self, host='localhost', port=3001):
        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server.bind((host, port))
        self.server.listen(1)
        self.clients = []
        self.running = True
    
    def accept_clients(self):
        """接受客户端连接"""
        while self.running:
            try:
                client, addr = self.server.accept()
                self.clients.append(client)
                threading.Thread(
                    target=self.handle_client,
                    args=(client, addr)
                ).start()
            except Exception as e:
                if self.running:
                    logger.error(f"接受客户端错误: {e}")
    
    def handle_client(self, client, addr):
        """处理客户端连接"""
        try:
            logger.info(f"新客户端连接: {addr}")
            while self.running:
                data = client.recv(4096)
                if not data:
                    break
                # 处理网络消息
                try:
                    request = json.loads(data.decode('utf-8'))
                    # 这里可以添加消息转发逻辑
                    logger.info(f"收到网络请求: {request.get('method')}")
                except Exception as e:
                    logger.error(f"处理网络消息错误: {e}")
        finally:
            client.close()
            logger.info(f"客户端断开: {addr}")

class NetworkMcpServer:
    """支持多线程和网络结构化的MCP服务器"""
    
    def __init__(self, config=None):
        # 解析配置
        self.config = config or {}
        transport_config = self.config.get("transport", {})
        network_config = transport_config.get("network", {})
        
        # 网络参数
        self.host = network_config.get("host", "localhost")
        self.port = network_config.get("port", 3001)
        self.timeout = network_config.get("timeout", 30)
        
        # 初始化组件
        self.thread_pool = ThreadPoolExecutor(
            max_workers=min(32, (os.cpu_count() or 1) * 4)
        )
        self.env_manager = EnvironmentManager()
        self.message_broker = MessageBroker()
        self.router = RequestRouter(self.thread_pool)
        
        # 请求ID计数器
        self._request_id = 0
        self._id_lock = threading.Lock()
        
        # 服务器状态
        self.running = True
        
        # 初始化网络服务器
        if network_config.get("enabled", True):
            self.network = NetworkServer(self.host, self.port)
        
        # 注册信号处理
        signal.signal(signal.SIGINT, self._handle_signal)
        signal.signal(signal.SIGTERM, self._handle_signal)
        
        # 注册路由
        self._register_routes()
        
        # 创建事件循环
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
    
    def _handle_signal(self, sig, frame):
        """处理信号"""
        logger.info("正在关闭服务器...")
        self.running = False
        self.thread_pool.shutdown(wait=True)
        sys.exit(0)
    
    def _register_routes(self):
        """注册默认路由"""
        self.router.register_route("list_tools", self._handle_list_tools)
        self.router.register_route("list_resources", self._handle_list_resources)
        self.router.register_route("call_tool", self._handle_call_tool)
        self.router.register_route("read_resource", self._handle_read_resource)
    
    def _get_next_id(self) -> int:
        """获取下一个请求ID"""
        with self._id_lock:
            self._request_id += 1
            return self._request_id
    
    async def _handle_list_tools(self, params: Dict) -> Dict:
        """处理列出工具请求"""
        tools = []
        
        # 从每个环境收集工具
        for env_name, env in self.env_manager.environments.items():
            with self.env_manager.use_environment(env_name):
                # 这里可以添加特定环境的工具
                if env_name == "python":
                    tools.extend([
                        {
                            "name": "async_echo",
                            "description": "异步回显工具",
                            "environment": "python",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "message": {
                                        "type": "string",
                                        "description": "要回显的消息"
                                    }
                                },
                                "required": ["message"]
                            }
                        },
                        {
                            "name": "system_info",
                            "description": "获取系统信息",
                            "environment": "python",
                            "inputSchema": {
                                "type": "object",
                                "properties": {}
                            }
                        },
                        {
                            "name": "pam_repair",
                            "description": "修复PAM配置问题",
                            "environment": "python",
                            "inputSchema": {
                                "type": "object",
                                "properties": {
                                    "action": {
                                        "type": "string",
                                        "enum": ["reinstall", "reset", "verify"],
                                        "description": "修复操作类型"
                                    }
                                },
                                "required": ["action"]
                            }
                        }
                    ])
        
        return {"tools": tools}
    
    async def _handle_list_resources(self, params: Dict) -> Dict:
        """处理列出资源请求"""
        resources = [
            {
                "uri": "system://info/resources",
                "name": "系统资源信息",
                "mimeType": "application/json",
                "description": "提供系统资源使用情况"
            }
        ]
        return {"resources": resources}
    
    async def _handle_call_tool(self, params: Dict) -> Dict:
        """处理调用工具请求"""
        tool_name = params.get("name")
        arguments = params.get("arguments", {})
        
        if tool_name == "async_echo":
            message = arguments.get("message", "")
            # 模拟异步处理
            await asyncio.sleep(0.1)
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"异步回显: {message}"
                    }
                ]
            }
        elif tool_name == "system_info":
            # 使用标准库获取系统信息
            import platform
            import os
            
            try:
                # 尝试导入psutil，如果可用
                import psutil
                has_psutil = True
            except ImportError:
                has_psutil = False
            
            if has_psutil:
                info = {
                    "cpu_percent": psutil.cpu_percent(interval=1),
                    "memory": dict(psutil.virtual_memory()._asdict()),
                    "disk": dict(psutil.disk_usage('/')._asdict()),
                    "network": len(psutil.net_connections())
                }
            else:
                # 使用标准库替代
                info = {
                    "system": platform.system(),
                    "platform": platform.platform(),
                    "python_version": platform.python_version(),
                    "cpu_count": os.cpu_count(),
                    "cwd": os.getcwd(),
                    "pid": os.getpid()
                }
                
                # 尝试获取内存信息
                try:
                    with open('/proc/meminfo', 'r') as f:
                        mem_info = {}
                        for line in f:
                            if ':' in line:
                                key, value = line.split(':', 1)
                                mem_info[key.strip()] = value.strip()
                        info["memory"] = mem_info
                except:
                    info["memory"] = "无法获取内存信息"
            
            return {
                "content": [
                    {
                        "type": "text",
                        "text": json.dumps(info, indent=2)
                    }
                ]
            }
        elif tool_name == "pam_repair":
            action = arguments.get("action", "verify")
            
            # 构建修复命令
            if action == "reinstall":
                cmd = ["pkexec", "apt", "install", "--reinstall", 
                      "pam", "libpam-modules", "sudo"]
            elif action == "reset":
                cmd = ["pkexec", "cp", "/etc/pam.d/sudo.dist", "/etc/pam.d/sudo"]
            else:  # verify
                cmd = ["pkexec", "diff", "/etc/pam.d/sudo", "/etc/pam.d/sudo.dist"]
            
            # 执行命令
            import subprocess
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True
            )
            
            if result.returncode != 0:
                return {
                    "content": [
                        {
                            "type": "text",
                            "text": f"PAM修复失败: {result.stderr}"
                        }
                    ],
                    "isError": True
                }
            
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"PAM修复成功({action}):\n{result.stdout}"
                    }
                ]
            }
        else:
            raise ValueError(f"未知工具: {tool_name}")
    
    async def _handle_read_resource(self, params: Dict) -> Dict:
        """处理读取资源请求"""
        uri = params.get("uri")
        
        if uri == "system://info/resources":
            import os
            import platform
            
            # 基本系统信息
            info = {
                "system": platform.system(),
                "platform": platform.platform(),
                "architecture": platform.machine(),
                "cpu_count": os.cpu_count(),
                "python_version": platform.python_version()
            }
            
            # 尝试获取更详细的系统信息
            try:
                # CPU信息
                with open('/proc/cpuinfo', 'r') as f:
                    cpu_info = {}
                    for line in f:
                        if ':' in line:
                            key, value = line.split(':', 1)
                            cpu_info[key.strip()] = value.strip()
                    info["cpu_info"] = cpu_info
            except:
                info["cpu_info"] = "无法获取CPU详细信息"
            
            # 内存信息
            try:
                with open('/proc/meminfo', 'r') as f:
                    mem_info = {}
                    for line in f:
                        if ':' in line:
                            key, value = line.split(':', 1)
                            mem_info[key.strip()] = value.strip()
                    info["memory_info"] = mem_info
            except:
                info["memory_info"] = "无法获取内存详细信息"
            
            # 磁盘信息
            try:
                import shutil
                total, used, free = shutil.disk_usage('/')
                info["disk_info"] = {
                    "total": total,
                    "used": used,
                    "free": free
                }
            except:
                info["disk_info"] = "无法获取磁盘信息"
            
            return {
                "contents": [
                    {
                        "uri": uri,
                        "mimeType": "application/json",
                        "text": json.dumps(info, indent=2)
                    }
                ]
            }
        else:
            raise ValueError(f"未知资源: {uri}")
    
    async def _process_request(self, request: Dict):
        """处理单个请求"""
        try:
            method = request.get("method")
            params = request.get("params", {})
            req_id = request.get("id", self._get_next_id())
            
            logger.info(f"处理请求 {req_id}: {method}")
            
            # 通过路由器处理请求
            result = await self.router.route_request(method, params)
            
            # 发送响应
            response = {
                "jsonrpc": "2.0",
                "id": req_id,
                "result": result
            }
            
            # 发布响应消息
            self.message_broker.publish("responses", response)
            
            # 写入响应
            sys.stdout.write(json.dumps(response) + "\n")
            sys.stdout.flush()
            
        except Exception as e:
            logger.error(f"处理请求时出错: {e}")
            error_response = {
                "jsonrpc": "2.0",
                "id": req_id,
                "error": {
                    "code": -32000,
                    "message": str(e)
                }
            }
            sys.stdout.write(json.dumps(error_response) + "\n")
            sys.stdout.flush()
    
    async def _input_loop(self):
        """处理输入的异步循环"""
        while self.running:
            try:
                # 非阻塞读取
                line = await self.loop.run_in_executor(
                    self.thread_pool,
                    sys.stdin.readline
                )
                
                if not line:
                    break
                
                try:
                    request = json.loads(line)
                    # 创建任务处理请求
                    self.loop.create_task(self._process_request(request))
                except json.JSONDecodeError:
                    logger.error("无效的JSON请求")
                    
            except Exception as e:
                logger.error(f"处理输入时出错: {e}")
    
    def _init_basic_components(self):
        """初始化基础组件"""
        # 初始化线程池
        self.thread_pool = ThreadPoolExecutor(max_workers=4)
        # 初始化事件循环
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        # 设置运行标志
        self.running = True

    def _load_base_topology(self):
        """加载基础网络拓扑"""
        # 初始化空网络结构
        self.nodes = []
        self.connections = []
        
        # 加载预设拓扑配置
        config_path = os.path.join(os.path.dirname(__file__), "network_config.json")
        if os.path.exists(config_path):
            with open(config_path) as f:
                config = json.load(f)
                self.nodes = config.get("nodes", [])
                self.connections = config.get("connections", [])
        
        logger.info(f"加载完成: {len(self.nodes)}个节点, {len(self.connections)}条连接")

    def _build_node_relations(self):
        """构建节点关系图"""
        self.relation_graph = defaultdict(list)
        for conn in self.connections:
            src, dst = conn["source"], conn["target"]
            self.relation_graph[src].append(dst)
            self.relation_graph[dst].append(src)  # 无向图
        
        logger.info(f"构建完成: {len(self.relation_graph)}个节点关系")

    def _init_routing_table(self):
        """初始化路由表"""
        self.routing_table = {}
        for node in self.nodes:
            self.routing_table[node["id"]] = {
                "ip": node["ip"],
                "neighbors": self.relation_graph.get(node["id"], [])
            }
        
        logger.info(f"路由表初始化完成: {len(self.routing_table)}个节点")

    def _init_network_structures(self):
        """分步初始化网络结构"""
        logger.info("步骤1: 加载基础网络拓扑...")
        self._load_base_topology()
        
        logger.info("步骤2: 构建节点关系...")
        self._build_node_relations()
        
        logger.info("步骤3: 初始化路由表...")
        self._init_routing_table()

    def _run_diagnostics(self):
        """运行预启动诊断"""
        from network_diagnostics import NetworkDiagnostics
        diag = NetworkDiagnostics(os.path.join(os.path.dirname(__file__), "network_config.json"))
        if not diag.run_full_check():
            logger.error("预启动诊断失败，请检查日志")
            return False
        return True

    def _init_phases(self):
        """分阶段初始化服务器"""
        logger.info("阶段1: 初始化基础组件...")
        self._init_basic_components()
        logger.info("基础组件初始化完成")

        # 设置超时信号处理
        def timeout_handler(signum, frame):
            logger.error("当前网络结构初始化进度:")
            logger.error(f"节点数: {len(getattr(self, 'nodes', []))}")
            logger.error(f"连接数: {len(getattr(self, 'connections', []))}")
            raise TimeoutError("服务器启动超时")
        
        self._original_signal = signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(30)  # 延长至30秒超时
        
        logger.info("阶段2: 初始化网络结构...")
        self._init_network_structures()
        signal.alarm(0)  # 取消超时
        logger.info("网络结构初始化完成")

    def _start_network_service(self):
        """启动网络服务"""
        logger.info("启动网络服务...")
        self.network_thread = threading.Thread(
            target=self.network.accept_clients,
            daemon=True
        )
        self.network_thread.start()
        logger.info(f"网络服务已启动，监听于 {self.network.server.getsockname()}")

    def _main_loop(self):
        """主事件循环"""
        try:
            self.loop.run_until_complete(self._input_loop())
        except KeyboardInterrupt:
            logger.info("收到中断信号")
        except Exception as e:
            logger.error(f"服务器运行错误: {str(e)}")
            raise
        finally:
            self._cleanup()

    def _cleanup(self):
        """清理资源"""
        signal.alarm(0)
        if hasattr(self, '_original_signal'):
            signal.signal(signal.SIGALRM, self._original_signal)
        self.running = False
        self.thread_pool.shutdown(wait=True)
        self.loop.close()
        logger.info("服务器已关闭")

    def run(self):
        """重构后的运行方法"""
        logger.info("启动网络结构化MCP服务器...")
        
        # 1. 初始化诊断
        if not self._run_diagnostics():
            return False
            
        # 2. 分阶段初始化
        try:
            self._init_phases()
        except TimeoutError as e:
            logger.error(f"网络结构初始化超时: {str(e)}")
            self._force_shutdown()
            return False
        except Exception as e:
            logger.error(f"初始化错误: {str(e)}")
            self._force_shutdown()
            return False
        
        # 3. 启动网络服务
        self._start_network_service()
        
        # 4. 主事件循环
        try:
            self._main_loop()
        except Exception as e:
            logger.error(f"服务器运行错误: {str(e)}")
            return False
        
        return True
    
    def _force_shutdown(self):
        """强制关闭服务器"""
        logger.warning("正在强制关闭服务器...")
        self.running = False
        try:
            self.thread_pool.shutdown(wait=False)
            self.loop.stop()
            self.loop.close()
        except:
            pass

if __name__ == "__main__":
    # 加载配置
    config_path = os.path.join(os.path.dirname(__file__), "python_mcp_config.json")
    try:
        with open(config_path) as f:
            config = json.load(f)
    except Exception as e:
        logger.error(f"加载配置文件失败: {str(e)}")
        config = {}
    
    # 启动服务器
    try:
        server = NetworkMcpServer(config)
        server.run()
    except Exception as e:
        logger.error(f"服务器启动失败: {str(e)}")
        sys.exit(1)